// In-memory storage implementation
use codegraph_core::{
    GraphStorage, Transaction, GraphQuery, QueryResult, CodeNode, CodeEdge,
    NodeId, EdgeId, EdgeType, Result, NodeType,
};
use async_trait::async_trait;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;
use serde::{Serialize, Deserialize};

/// Serializable node data for storage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredNode {
    pub id: NodeId,
    pub node_type: NodeType,
    pub name: String,
    pub location: codegraph_core::SourceLocation,
    pub language: String,
    pub metadata: codegraph_core::NodeMetadata,
}

/// Serializable edge data for storage
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StoredEdge {
    pub id: EdgeId,
    pub source: NodeId,
    pub target: NodeId,
    pub edge_type: EdgeType,
    pub metadata: codegraph_core::EdgeMetadata,
}

pub struct MemoryStorage {
    nodes: Arc<RwLock<HashMap<NodeId, StoredNode>>>,
    edges: Arc<RwLock<HashMap<EdgeId, StoredEdge>>>,
    node_edges: Arc<RwLock<HashMap<NodeId, Vec<EdgeId>>>>, // Node -> its edges
}

impl MemoryStorage {
    pub fn new() -> Self {
        Self {
            nodes: Arc::new(RwLock::new(HashMap::new())),
            edges: Arc::new(RwLock::new(HashMap::new())),
            node_edges: Arc::new(RwLock::new(HashMap::new())),
        }
    }
}

impl Default for MemoryStorage {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl GraphStorage for MemoryStorage {
    async fn save_node(&self, node: &dyn CodeNode) -> Result<()> {
        let stored_node = StoredNode {
            id: node.id(),
            node_type: node.node_type().clone(),
            name: node.name().to_string(),
            location: node.location().clone(),
            language: node.language().to_string(),
            metadata: node.metadata().clone(),
        };

        let mut nodes = self.nodes.write().await;
        nodes.insert(node.id(), stored_node);

        Ok(())
    }

    async fn get_node(&self, id: NodeId) -> Result<Option<Box<dyn CodeNode>>> {
        let nodes = self.nodes.read().await;
        if let Some(stored_node) = nodes.get(&id) {
            // Convert StoredNode back to a CodeNode implementation
            let node = codegraph_graph::GraphNode::new(
                stored_node.id,
                stored_node.node_type.clone(),
                stored_node.name.clone(),
                stored_node.location.clone(),
                stored_node.language.clone(),
            );
            Ok(Some(Box::new(node)))
        } else {
            Ok(None)
        }
    }

    async fn delete_node(&self, id: NodeId) -> Result<bool> {
        let mut nodes = self.nodes.write().await;
        let existed = nodes.remove(&id).is_some();

        if existed {
            // Also remove all edges connected to this node
            let mut edges = self.edges.write().await;
            let mut node_edges = self.node_edges.write().await;

            // Remove edges where this node is source or target
            edges.retain(|_, edge| edge.source != id && edge.target != id);

            // Remove from node_edges mapping
            node_edges.remove(&id);

            // Remove this node from other nodes' edge lists
            for edge_list in node_edges.values_mut() {
                edge_list.retain(|edge_id| {
                    if let Some(edge) = edges.get(edge_id) {
                        edge.source != id && edge.target != id
                    } else {
                        false
                    }
                });
            }
        }

        Ok(existed)
    }
    
    async fn save_edge(&self, edge: &dyn CodeEdge) -> Result<()> {
        let stored_edge = StoredEdge {
            id: edge.id(),
            source: edge.source(),
            target: edge.target(),
            edge_type: edge.edge_type().clone(),
            metadata: edge.metadata().clone(),
        };

        let mut edges = self.edges.write().await;
        let mut node_edges = self.node_edges.write().await;

        // Store the edge
        edges.insert(edge.id(), stored_edge);

        // Update node-to-edges mapping
        node_edges.entry(edge.source()).or_insert_with(Vec::new).push(edge.id());
        node_edges.entry(edge.target()).or_insert_with(Vec::new).push(edge.id());

        Ok(())
    }

    async fn get_edges(&self, node_id: NodeId, edge_type: Option<&EdgeType>) -> Result<Vec<Box<dyn CodeEdge>>> {
        let edges = self.edges.read().await;
        let node_edges = self.node_edges.read().await;

        let mut result = Vec::new();

        if let Some(edge_ids) = node_edges.get(&node_id) {
            for edge_id in edge_ids {
                if let Some(stored_edge) = edges.get(edge_id) {
                    // Filter by edge type if specified
                    if let Some(filter_type) = edge_type {
                        if &stored_edge.edge_type != filter_type {
                            continue;
                        }
                    }

                    // Convert StoredEdge back to a CodeEdge implementation
                    let edge = codegraph_graph::GraphEdge::new(
                        stored_edge.id,
                        stored_edge.source,
                        stored_edge.target,
                        stored_edge.edge_type.clone(),
                    );
                    result.push(Box::new(edge) as Box<dyn CodeEdge>);
                }
            }
        }

        Ok(result)
    }

    async fn delete_edge(&self, id: EdgeId) -> Result<bool> {
        let mut edges = self.edges.write().await;
        let mut node_edges = self.node_edges.write().await;

        if let Some(stored_edge) = edges.remove(&id) {
            // Remove from node-to-edges mapping
            if let Some(source_edges) = node_edges.get_mut(&stored_edge.source) {
                source_edges.retain(|&edge_id| edge_id != id);
            }
            if let Some(target_edges) = node_edges.get_mut(&stored_edge.target) {
                target_edges.retain(|&edge_id| edge_id != id);
            }
            Ok(true)
        } else {
            Ok(false)
        }
    }
    
    async fn save_batch(
        &self,
        nodes: &[&dyn CodeNode],
        edges: &[&dyn CodeEdge],
    ) -> Result<()> {
        // Save all nodes first
        for node in nodes {
            self.save_node(*node).await?;
        }

        // Then save all edges
        for edge in edges {
            self.save_edge(*edge).await?;
        }

        Ok(())
    }

    async fn query(&self, query: &GraphQuery) -> Result<QueryResult> {
        let start_time = std::time::Instant::now();

        let nodes_guard = self.nodes.read().await;
        let edges_guard = self.edges.read().await;

        let mut result_node_ids = Vec::new();
        let mut result_edge_ids = Vec::new();

        // Filter nodes based on query criteria
        for (_, stored_node) in nodes_guard.iter() {
            let mut matches = true;

            // Filter by node types
            if let Some(ref node_types) = query.node_types {
                if !node_types.contains(&stored_node.node_type) {
                    matches = false;
                }
            }

            // Filter by languages
            if let Some(ref languages) = query.languages {
                if !languages.contains(&stored_node.language) {
                    matches = false;
                }
            }

            // Filter by name pattern (simple contains check)
            if let Some(ref pattern) = query.name_pattern {
                if !stored_node.name.contains(pattern) {
                    matches = false;
                }
            }

            if matches {
                result_node_ids.push(stored_node.id);
            }
        }

        // Filter edges based on query criteria
        for (_, stored_edge) in edges_guard.iter() {
            let mut matches = true;

            // Filter by edge types
            if let Some(ref edge_types) = query.edge_types {
                if !edge_types.contains(&stored_edge.edge_type) {
                    matches = false;
                }
            }

            if matches {
                result_edge_ids.push(stored_edge.id);
            }
        }

        // Apply limit and offset
        if let Some(offset) = query.offset {
            if offset < result_node_ids.len() {
                result_node_ids = result_node_ids.into_iter().skip(offset).collect();
            } else {
                result_node_ids.clear();
            }
        }

        if let Some(limit) = query.limit {
            result_node_ids.truncate(limit);
        }

        let execution_time = start_time.elapsed().as_millis() as u64;

        Ok(QueryResult {
            nodes: result_node_ids,
            edges: result_edge_ids,
            metadata: std::collections::HashMap::new(),
            execution_time_ms: execution_time,
        })
    }
    
    async fn begin_transaction(&self) -> Result<Box<dyn Transaction>> {
        Ok(Box::new(MemoryTransaction::new(
            self.nodes.clone(),
            self.edges.clone(),
            self.node_edges.clone(),
        )))
    }
}

pub struct MemoryTransaction {
    nodes: Arc<RwLock<HashMap<NodeId, StoredNode>>>,
    edges: Arc<RwLock<HashMap<EdgeId, StoredEdge>>>,
    node_edges: Arc<RwLock<HashMap<NodeId, Vec<EdgeId>>>>,
    // Track changes for rollback - using RwLock for interior mutability
    pending_nodes: Arc<RwLock<HashMap<NodeId, StoredNode>>>,
    pending_edges: Arc<RwLock<HashMap<EdgeId, StoredEdge>>>,
    deleted_nodes: Arc<RwLock<Vec<NodeId>>>,
    deleted_edges: Arc<RwLock<Vec<EdgeId>>>,
}

impl MemoryTransaction {
    pub fn new(
        nodes: Arc<RwLock<HashMap<NodeId, StoredNode>>>,
        edges: Arc<RwLock<HashMap<EdgeId, StoredEdge>>>,
        node_edges: Arc<RwLock<HashMap<NodeId, Vec<EdgeId>>>>,
    ) -> Self {
        Self {
            nodes,
            edges,
            node_edges,
            pending_nodes: Arc::new(RwLock::new(HashMap::new())),
            pending_edges: Arc::new(RwLock::new(HashMap::new())),
            deleted_nodes: Arc::new(RwLock::new(Vec::new())),
            deleted_edges: Arc::new(RwLock::new(Vec::new())),
        }
    }
}

#[async_trait]
impl Transaction for MemoryTransaction {
    async fn commit(self: Box<Self>) -> Result<()> {
        // Apply all pending changes to the main storage
        let mut nodes = self.nodes.write().await;
        let mut edges = self.edges.write().await;
        let mut node_edges = self.node_edges.write().await;

        // Apply node changes
        let mut pending_nodes = self.pending_nodes.write().await;
        for (id, node) in pending_nodes.drain() {
            nodes.insert(id, node);
        }

        // Apply edge changes
        let mut pending_edges = self.pending_edges.write().await;
        for (id, edge) in pending_edges.drain() {
            edges.insert(id, edge.clone());

            // Update node-to-edges mapping
            node_edges.entry(edge.source).or_insert_with(Vec::new).push(id);
            node_edges.entry(edge.target).or_insert_with(Vec::new).push(id);
        }

        // Apply deletions
        let deleted_nodes = self.deleted_nodes.read().await;
        for node_id in deleted_nodes.iter() {
            nodes.remove(node_id);
            node_edges.remove(node_id);
        }

        let deleted_edges = self.deleted_edges.read().await;
        for edge_id in deleted_edges.iter() {
            if let Some(edge) = edges.remove(edge_id) {
                // Remove from node-to-edges mapping
                if let Some(source_edges) = node_edges.get_mut(&edge.source) {
                    source_edges.retain(|&id| id != *edge_id);
                }
                if let Some(target_edges) = node_edges.get_mut(&edge.target) {
                    target_edges.retain(|&id| id != *edge_id);
                }
            }
        }

        Ok(())
    }

    async fn rollback(self: Box<Self>) -> Result<()> {
        // Simply drop the transaction - all pending changes are discarded
        Ok(())
    }

    async fn save_node(&self, node: &dyn CodeNode) -> Result<()> {
        let stored_node = StoredNode {
            id: node.id(),
            node_type: node.node_type().clone(),
            name: node.name().to_string(),
            location: node.location().clone(),
            language: node.language().to_string(),
            metadata: node.metadata().clone(),
        };

        let mut pending_nodes = self.pending_nodes.write().await;
        pending_nodes.insert(node.id(), stored_node);
        Ok(())
    }

    async fn save_edge(&self, edge: &dyn CodeEdge) -> Result<()> {
        let stored_edge = StoredEdge {
            id: edge.id(),
            source: edge.source(),
            target: edge.target(),
            edge_type: edge.edge_type().clone(),
            metadata: edge.metadata().clone(),
        };

        let mut pending_edges = self.pending_edges.write().await;
        pending_edges.insert(edge.id(), stored_edge);
        Ok(())
    }

    async fn delete_node(&self, id: NodeId) -> Result<bool> {
        // Check if node exists (either in main storage or pending)
        let pending_nodes = self.pending_nodes.read().await;
        let exists = pending_nodes.contains_key(&id) || {
            let nodes = self.nodes.read().await;
            nodes.contains_key(&id)
        };

        if exists {
            let mut deleted_nodes = self.deleted_nodes.write().await;
            deleted_nodes.push(id);
            // Also remove from pending if it was there
            drop(pending_nodes);
            let mut pending_nodes = self.pending_nodes.write().await;
            pending_nodes.remove(&id);
        }

        Ok(exists)
    }

    async fn delete_edge(&self, id: EdgeId) -> Result<bool> {
        // Check if edge exists (either in main storage or pending)
        let pending_edges = self.pending_edges.read().await;
        let exists = pending_edges.contains_key(&id) || {
            let edges = self.edges.read().await;
            edges.contains_key(&id)
        };

        if exists {
            let mut deleted_edges = self.deleted_edges.write().await;
            deleted_edges.push(id);
            // Also remove from pending if it was there
            drop(pending_edges);
            let mut pending_edges = self.pending_edges.write().await;
            pending_edges.remove(&id);
        }

        Ok(exists)
    }
}
